import pymysql
from dotenv import load_dotenv
import os
import logging
from pymysql import Error
from pymysql.cursors import DictCursor

load_dotenv()

class Database:
    def __init__(self):
        self.connection = None
        self.connect()

    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = pymysql.connect(
                host=os.getenv('MYSQL_HOST'),
                port=int(os.getenv('MYSQL_PORT')),
                user=os.getenv('MYSQL_USER'),
                password=os.getenv('MYSQL_PASSWORD'),
                database=os.getenv('MYSQL_DATABASE'),
                charset='utf8mb4',
                cursorclass=DictCursor,
                autocommit=True
            )
            logging.info("数据库连接成功")
        except Error as e:
            logging.error(f"数据库连接失败: {str(e)}")
            raise

    def execute_query(self, query):
        """执行SQL查询"""
        if not self.connection:
            self.connect()
            
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(query)
                
                # 处理DESCRIBE命令
                if query.strip().upper().startswith('DESCRIBE'):
                    result = cursor.fetchall()
                    # 将结果转换为字典列表
                    formatted_result = []
                    for row in result:
                        formatted_result.append({
                            "Field": row["Field"],
                            "Type": row["Type"],
                            "Null": row["Null"],
                            "Key": row["Key"],
                            "Default": row["Default"],
                            "Extra": row["Extra"]
                        })
                    return formatted_result
                
                # 处理SHOW TABLES等特殊命令
                elif query.strip().upper().startswith('SHOW'):
                    result = cursor.fetchall()
                    if isinstance(result[0], dict):
                        return result
                    else:
                        return [{"Tables_in_" + os.getenv('MYSQL_DATABASE'): row[0]} for row in result]
                
                # 处理SELECT查询
                elif query.strip().upper().startswith('SELECT'):
                    result = cursor.fetchall()
                    logging.info(f"查询成功，返回 {len(result)} 条记录")
                    return result
                
                # 处理其他操作（INSERT, UPDATE, DELETE等）
                else:
                    affected_rows = cursor.rowcount
                    logging.info(f"操作成功，影响 {affected_rows} 行")
                    return {
                        "status": "success",
                        "affected_rows": affected_rows,
                        "message": f"操作成功，影响 {affected_rows} 行"
                    }
        except Error as e:
            error_msg = f"SQL执行错误: {str(e)}"
            logging.error(error_msg)
            return {
                "status": "error",
                "message": error_msg
            }

    def close(self):
        """关闭数据库连接"""
        if self.connection:
            try:
                self.connection.close()
                logging.info("数据库连接已关闭")
            except Error as e:
                logging.error(f"关闭数据库连接时发生错误: {str(e)}")
            finally:
                self.connection = None 